import pygame
import random
import time
import sys
from enum import Enum

# 初始化pygame
pygame.init()

# 颜色定义
class Colors:
    BLACK = (0, 0, 0)
    WHITE = (255, 255, 255)
    RED = (255, 0, 0)
    GREEN = (0, 255, 0)
    BLUE = (0, 0, 255)
    GRAY = (128, 128, 128)
    YELLOW = (255, 255, 0)
    PURPLE = (128, 0, 128)

# 方向枚举
class Direction(Enum):
    UP = 0
    DOWN = 1
    LEFT = 2
    RIGHT = 3

# 游戏难度设置
class Difficulty:
    EASY = {"speed": 10, "initial_length": 5, "food_size": 20}
    NORMAL = {"speed": 15, "initial_length": 3, "food_size": 15}
    HARD = {"speed": 20, "initial_length": 1, "food_size": 10}

# 食物类型
class FoodType(Enum):
    NORMAL = 0  # 普通食物
    SPEED_UP = 1  # 加速食物
    SPEED_DOWN = 2  # 减速食物
    INVINCIBLE = 3  # 无敌食物
    SCORE_DOWN = 4  # 减分食物

# 游戏设置类
class GameSettings:
    def __init__(self):
        self.difficulty = Difficulty.NORMAL
        self.snake_color = Colors.GREEN
        self.food_color = Colors.RED
        self.bg_color = Colors.BLACK
        self.show_grid = True
        self.boundary_mode = "death"  # "death" or "pass_through"
        self.enable_special_food = False
        self.enable_obstacles = False

# 食物类
class Food:
    def __init__(self, x, y, size, food_type=FoodType.NORMAL):
        self.x = x
        self.y = y
        self.size = size
        self.type = food_type
        self.eaten = False
        
        # 根据食物类型设置颜色
        if self.type == FoodType.NORMAL:
            self.color = Colors.RED
        elif self.type == FoodType.SPEED_UP:
            self.color = Colors.YELLOW
        elif self.type == FoodType.SPEED_DOWN:
            self.color = Colors.BLUE
        elif self.type == FoodType.INVINCIBLE:
            self.color = Colors.PURPLE
        elif self.type == FoodType.SCORE_DOWN:
            self.color = Colors.GRAY
    
    def draw(self, surface):
        pygame.draw.rect(surface, self.color, [self.x, self.y, self.size, self.size])

# 蛇类
class Snake:
    def __init__(self, x, y, block_size, initial_length=3):
        self.blocks = []
        self.block_size = block_size
        self.direction = Direction.RIGHT
        self.next_direction = Direction.RIGHT
        self.speed = 15  # 初始速度
        self.invincible = False  # 无敌状态
        self.invincible_end_time = 0
        
        # 初始化蛇身
        for i in range(initial_length):
            self.blocks.append([x - i * block_size, y])
    
    def change_direction(self, new_direction):
        # 防止180度急转弯
        if (new_direction == Direction.UP and self.direction != Direction.DOWN) or \
           (new_direction == Direction.DOWN and self.direction != Direction.UP) or \
           (new_direction == Direction.LEFT and self.direction != Direction.RIGHT) or \
           (new_direction == Direction.RIGHT and self.direction != Direction.LEFT):
            self.next_direction = new_direction
    
    def move(self):
        # 更新方向
        self.direction = self.next_direction
        
        # 获取蛇头位置
        head = self.blocks[0].copy()
        
        # 根据方向移动蛇头
        if self.direction == Direction.UP:
            head[1] -= self.block_size
        elif self.direction == Direction.DOWN:
            head[1] += self.block_size
        elif self.direction == Direction.LEFT:
            head[0] -= self.block_size
        elif self.direction == Direction.RIGHT:
            head[0] += self.block_size
        
        # 在列表开头插入新的蛇头
        self.blocks.insert(0, head)
        
        # 移除蛇尾（除非吃到食物）
        return self.blocks.pop()
    
    def grow(self):
        # 蛇增长时，保留蛇尾
        tail = self.blocks[-1]
        self.blocks.append(tail.copy())
    
    def draw(self, surface, color):
        for i, block in enumerate(self.blocks):
            # 蛇头用不同颜色
            block_color = Colors.YELLOW if i == 0 else color
            pygame.draw.rect(surface, block_color, [block[0], block[1], self.block_size, self.block_size])
    
    def check_collision(self, width, height, boundary_mode):
        head = self.blocks[0]
        
        # 检查是否撞墙
        if boundary_mode == "death":
            if head[0] < 0 or head[0] >= width or head[1] < 0 or head[1] >= height:
                return True
        
        # 检查是否撞到自己（无敌状态下不会）
        if not self.invincible:
            for block in self.blocks[1:]:
                if head[0] == block[0] and head[1] == block[1]:
                    return True
        
        # 检查无敌状态是否结束
        if self.invincible and time.time() > self.invincible_end_time:
            self.invincible = False
        
        return False
    
    def check_boundary_cross(self, width, height):
        head = self.blocks[0]
        new_head = head.copy()
        
        if head[0] < 0:
            new_head[0] = width - self.block_size
        elif head[0] >= width:
            new_head[0] = 0
        elif head[1] < 0:
            new_head[1] = height - self.block_size
        elif head[1] >= height:
            new_head[1] = 0
        
        if new_head != head:
            self.blocks[0] = new_head
            return True
        return False

# 障碍物类
class Obstacle:
    def __init__(self, x, y, width, height):
        self.x = x
        self.y = y
        self.width = width
        self.height = height
        self.color = Colors.GRAY
    
    def draw(self, surface):
        pygame.draw.rect(surface, self.color, [self.x, self.y, self.width, self.height])

# 游戏类
class SnakeGame:
    def __init__(self, width=800, height=600):
        self.width = width
        self.height = height
        self.block_size = 20
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption('贪吃蛇游戏')
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font("C:/Windows/Fonts/simhei.ttf", 36)
        self.small_font = pygame.font.Font("C:/Windows/Fonts/simhei.ttf", 24)
        self.settings = GameSettings()
        self.reset_game()
        self.state = "menu"  # menu, game, game_over, settings
        self.high_score = 0
        self.load_high_score()
    
    def reset_game(self):
        # 初始化蛇
        initial_length = self.settings.difficulty["initial_length"]
        start_x = self.width // 2
        start_y = self.height // 2
        self.snake = Snake(start_x, start_y, self.block_size, initial_length)
        self.snake.speed = self.settings.difficulty["speed"]
        
        # 初始化食物
        self.foods = []
        self.spawn_food()
        
        # 初始化障碍物
        self.obstacles = []
        if self.settings.enable_obstacles:
            self.spawn_obstacles()
        
        # 游戏状态
        self.score = 0
        self.game_over = False
        self.paused = False
        self.combo = 0
        self.last_eat_time = 0
        self.combo_end_time = 0
    
    def spawn_food(self):
        # 生成食物位置，确保不与蛇身重叠
        while True:
            x = random.randint(0, (self.width - self.block_size) // self.block_size) * self.block_size
            y = random.randint(0, (self.height - self.block_size) // self.block_size) * self.block_size
            
            # 检查是否与蛇身重叠
            overlap = False
            for block in self.snake.blocks:
                if block[0] == x and block[1] == y:
                    overlap = True
                    break
            
            # 检查是否与障碍物重叠
            if not overlap and self.settings.enable_obstacles:
                for obstacle in self.obstacles:
                    if (x >= obstacle.x and x < obstacle.x + obstacle.width and
                        y >= obstacle.y and y < obstacle.y + obstacle.height):
                        overlap = True
                        break
            
            if not overlap:
                break
        
        # 决定食物类型（5%概率生成特殊食物）
        food_type = FoodType.NORMAL
        if self.settings.enable_special_food and random.random() < 0.05:
            food_type = random.choice(list(FoodType))
        
        food_size = self.settings.difficulty["food_size"]
        self.foods.append(Food(x, y, food_size, food_type))
    
    def spawn_obstacles(self):
        # 生成固定障碍物
        for _ in range(5):
            width = random.randint(1, 3) * self.block_size
            height = random.randint(1, 3) * self.block_size
            x = random.randint(0, (self.width - width) // self.block_size) * self.block_size
            y = random.randint(0, (self.height - height) // self.block_size) * self.block_size
            self.obstacles.append(Obstacle(x, y, width, height))
    
    def check_food_collision(self):
        head = self.snake.blocks[0]
        
        for food in self.foods[:]:
            if (head[0] < food.x + food.size and head[0] + self.block_size > food.x and
                head[1] < food.y + food.size and head[1] + self.block_size > food.y):
                
                # 处理食物效果
                self.handle_food_effect(food)
                
                # 移除被吃掉的食物
                self.foods.remove(food)
                food.eaten = True
                
                # 蛇增长
                self.snake.grow()
                
                # 生成新食物
                self.spawn_food()
                
                # 更新分数
                self.update_score(food)
                
                return True
        
        return False
    
    def handle_food_effect(self, food):
        current_time = time.time()
        
        if food.type == FoodType.SPEED_UP:
            self.snake.speed += 5  # 加速
            # 5秒后恢复原速
            pygame.time.set_timer(pygame.USEREVENT, 5000)
        elif food.type == FoodType.SPEED_DOWN:
            self.snake.speed = max(5, self.snake.speed - 5)  # 减速，最低速度为5
            # 5秒后恢复原速
            pygame.time.set_timer(pygame.USEREVENT, 5000)
        elif food.type == FoodType.INVINCIBLE:
            self.snake.invincible = True
            self.snake.invincible_end_time = current_time + 10  # 10秒无敌时间
        elif food.type == FoodType.SCORE_DOWN:
            pass  # 分数更新在update_score方法中处理
    
    def update_score(self, food):
        base_score = 10
        
        # 根据食物类型调整分数
        if food.type == FoodType.NORMAL:
            self.score += base_score
        elif food.type == FoodType.SPEED_UP:
            self.score += base_score * 2
        elif food.type == FoodType.SPEED_DOWN:
            self.score += base_score // 2
        elif food.type == FoodType.INVINCIBLE:
            self.score += base_score * 3
        elif food.type == FoodType.SCORE_DOWN:
            self.score = max(0, self.score - base_score * 2)
        
        # 连击系统
        current_time = time.time()
        if current_time - self.last_eat_time < 1.0:  # 1秒内连续吃到食物
            self.combo += 1
            self.score += self.combo * 2  # 连击加分
            self.combo_end_time = current_time + 1.0  # 重置连击时间
        else:
            self.combo = 0
        
        self.last_eat_time = current_time
        
        # 更新最高分
        if self.score > self.high_score:
            self.high_score = self.score
            self.save_high_score()
    
    def check_obstacle_collision(self):
        if not self.settings.enable_obstacles or self.snake.invincible:
            return False
        
        head = self.snake.blocks[0]
        
        for obstacle in self.obstacles:
            if (head[0] < obstacle.x + obstacle.width and
                head[0] + self.block_size > obstacle.x and
                head[1] < obstacle.y + obstacle.height and
                head[1] + self.block_size > obstacle.y):
                return True
        
        return False
    
    def draw_grid(self):
        for x in range(0, self.width, self.block_size):
            pygame.draw.line(self.screen, Colors.GRAY, (x, 0), (x, self.height))
        for y in range(0, self.height, self.block_size):
            pygame.draw.line(self.screen, Colors.GRAY, (0, y), (self.width, y))
    
    def draw_score(self):
        score_text = self.font.render(f"分数: {self.score}", True, Colors.WHITE)
        high_score_text = self.font.render(f"最高分: {self.high_score}", True, Colors.WHITE)
        self.screen.blit(score_text, (10, 10))
        self.screen.blit(high_score_text, (10, 50))
        
        # 显示连击
        if self.combo > 0:
            combo_text = self.small_font.render(f"连击: {self.combo}", True, Colors.YELLOW)
            self.screen.blit(combo_text, (10, 90))
        
        # 显示无敌状态
        if self.snake.invincible:
            invincible_text = self.small_font.render("无敌!", True, Colors.PURPLE)
            self.screen.blit(invincible_text, (10, 120))
    
    def draw_pause(self):
        if self.paused:
            pause_text = self.font.render("游戏暂停 - 按P继续", True, Colors.WHITE)
            text_rect = pause_text.get_rect(center=(self.width//2, self.height//2))
            self.screen.blit(pause_text, text_rect)
    
    def draw_game_over(self):
        game_over_text = self.font.render("游戏结束!", True, Colors.RED)
        restart_text = self.font.render("按R重新开始", True, Colors.WHITE)
        menu_text = self.font.render("按M返回菜单", True, Colors.WHITE)
        
        self.screen.blit(game_over_text, (self.width//2 - game_over_text.get_width()//2, self.height//2 - 50))
        self.screen.blit(restart_text, (self.width//2 - restart_text.get_width()//2, self.height//2))
        self.screen.blit(menu_text, (self.width//2 - menu_text.get_width()//2, self.height//2 + 50))
        
        # 显示是否打破记录
        if self.score == self.high_score and self.score > 0:
            record_text = self.font.render("新纪录!", True, Colors.YELLOW)
            self.screen.blit(record_text, (self.width//2 - record_text.get_width()//2, self.height//2 + 100))
    
    def draw_menu(self):
        title_text = self.font.render("贪吃蛇游戏", True, Colors.GREEN)
        start_text = self.font.render("1. 开始游戏", True, Colors.WHITE)
        settings_text = self.font.render("2. 游戏设置", True, Colors.WHITE)
        quit_text = self.font.render("3. 退出游戏", True, Colors.WHITE)
        
        self.screen.blit(title_text, (self.width//2 - title_text.get_width()//2, self.height//4))
        self.screen.blit(start_text, (self.width//2 - start_text.get_width()//2, self.height//2))
        self.screen.blit(settings_text, (self.width//2 - settings_text.get_width()//2, self.height//2 + 50))
        self.screen.blit(quit_text, (self.width//2 - quit_text.get_width()//2, self.height//2 + 100))
        
        # 显示最高分
        high_score_text = self.font.render(f"最高分: {self.high_score}", True, Colors.YELLOW)
        self.screen.blit(high_score_text, (self.width//2 - high_score_text.get_width()//2, self.height//4 + 50))
    
    def draw_settings(self):
        title_text = self.font.render("游戏设置", True, Colors.GREEN)
        difficulty_text = self.font.render(f"难度: {self.get_difficulty_name()}", True, Colors.WHITE)
        boundary_text = self.font.render(f"边界模式: {'死亡' if self.settings.boundary_mode == 'death' else '穿墙'}", True, Colors.WHITE)
        special_food_text = self.font.render(f"特殊食物: {'开启' if self.settings.enable_special_food else '关闭'}", True, Colors.WHITE)
        obstacles_text = self.font.render(f"障碍物: {'开启' if self.settings.enable_obstacles else '关闭'}", True, Colors.WHITE)
        back_text = self.font.render("按B返回菜单", True, Colors.WHITE)
        
        self.screen.blit(title_text, (self.width//2 - title_text.get_width()//2, self.height//4))
        self.screen.blit(difficulty_text, (self.width//2 - difficulty_text.get_width()//2, self.height//2))
        self.screen.blit(boundary_text, (self.width//2 - boundary_text.get_width()//2, self.height//2 + 50))
        self.screen.blit(special_food_text, (self.width//2 - special_food_text.get_width()//2, self.height//2 + 100))
        self.screen.blit(obstacles_text, (self.width//2 - obstacles_text.get_width()//2, self.height//2 + 150))
        self.screen.blit(back_text, (self.width//2 - back_text.get_width()//2, self.height - 100))
    
    def get_difficulty_name(self):
        if self.settings.difficulty == Difficulty.EASY:
            return "简单"
        elif self.settings.difficulty == Difficulty.NORMAL:
            return "普通"
        else:
            return "困难"
    
    def toggle_difficulty(self):
        if self.settings.difficulty == Difficulty.EASY:
            self.settings.difficulty = Difficulty.NORMAL
        elif self.settings.difficulty == Difficulty.NORMAL:
            self.settings.difficulty = Difficulty.HARD
        else:
            self.settings.difficulty = Difficulty.EASY
    
    def toggle_boundary_mode(self):
        if self.settings.boundary_mode == "death":
            self.settings.boundary_mode = "pass_through"
        else:
            self.settings.boundary_mode = "death"
    
    def toggle_special_food(self):
        self.settings.enable_special_food = not self.settings.enable_special_food
    
    def toggle_obstacles(self):
        self.settings.enable_obstacles = not self.settings.enable_obstacles
    
    def save_high_score(self):
        try:
            with open("highscore.txt", "w") as f:
                f.write(str(self.high_score))
        except:
            pass
    
    def load_high_score(self):
        try:
            with open("highscore.txt", "r") as f:
                self.high_score = int(f.read())
        except:
            self.high_score = 0
    
    def run(self):
        running = True
        
        while running:
            # 处理事件
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                
                elif event.type == pygame.KEYDOWN:
                    if self.state == "game" and not self.game_over:
                        # 游戏中的按键
                        if event.key == pygame.K_UP:
                            self.snake.change_direction(Direction.UP)
                        elif event.key == pygame.K_DOWN:
                            self.snake.change_direction(Direction.DOWN)
                        elif event.key == pygame.K_LEFT:
                            self.snake.change_direction(Direction.LEFT)
                        elif event.key == pygame.K_RIGHT:
                            self.snake.change_direction(Direction.RIGHT)
                        elif event.key == pygame.K_p:
                            self.paused = not self.paused
                        elif event.key == pygame.K_m:
                            self.state = "menu"
                    
                    elif self.state == "game_over":
                        # 游戏结束后的按键
                        if event.key == pygame.K_r:
                            self.reset_game()
                            self.state = "game"
                        elif event.key == pygame.K_m:
                            self.state = "menu"
                    
                    elif self.state == "menu":
                        # 菜单中的按键
                        if event.key == pygame.K_1:
                            self.reset_game()
                            self.state = "game"
                        elif event.key == pygame.K_2:
                            self.state = "settings"
                        elif event.key == pygame.K_3:
                            running = False
                    
                    elif self.state == "settings":
                        # 设置中的按键
                        if event.key == pygame.K_1:
                            self.toggle_difficulty()
                        elif event.key == pygame.K_2:
                            self.toggle_boundary_mode()
                        elif event.key == pygame.K_3:
                            self.toggle_special_food()
                        elif event.key == pygame.K_4:
                            self.toggle_obstacles()
                        elif event.key == pygame.K_b:
                            self.state = "menu"
                
                elif event.type == pygame.USEREVENT:
                    # 恢复蛇的原始速度
                    self.snake.speed = self.settings.difficulty["speed"]
            
            # 清屏
            self.screen.fill(self.settings.bg_color)
            
            # 根据游戏状态绘制不同界面
            if self.state == "menu":
                self.draw_menu()
            
            elif self.state == "settings":
                self.draw_settings()
            
            elif self.state == "game":
                # 绘制网格
                if self.settings.show_grid:
                    self.draw_grid()
                
                # 绘制食物
                for food in self.foods:
                    food.draw(self.screen)
                
                # 绘制障碍物
                for obstacle in self.obstacles:
                    obstacle.draw(self.screen)
                
                # 绘制蛇
                self.snake.draw(self.screen, self.settings.snake_color)
                
                # 绘制分数
                self.draw_score()
                
                # 绘制暂停状态
                if self.paused:
                    self.draw_pause()
                elif not self.game_over:
                    # 游戏逻辑更新
                    self.update_game()
                
                # 绘制游戏结束
                if self.game_over:
                    self.draw_game_over()
            
            # 更新屏幕
            pygame.display.flip()
            
            # 控制帧率
            self.clock.tick(self.snake.speed)
        
        pygame.quit()
        sys.exit()
    
    def update_game(self):
        # 移动蛇
        self.snake.move()
        
        # 检查是否吃到食物
        self.check_food_collision()
        
        # 检查是否撞到障碍物
        if self.check_obstacle_collision():
            self.game_over = True
        
        # 检查边界碰撞
        if self.settings.boundary_mode == "death":
            if self.snake.check_collision(self.width, self.height, self.settings.boundary_mode):
                self.game_over = True
        else:
            self.snake.check_boundary_cross(self.width, self.height)
        
        # 检查是否撞到自己
        if self.snake.check_collision(self.width, self.height, self.settings.boundary_mode):
            self.game_over = True
        
        # 检查连击是否结束
        if self.combo > 0 and time.time() > self.combo_end_time:
            self.combo = 0

# 主函数
def main():
    game = SnakeGame()
    game.run()

if __name__ == "__main__":
    main()